package com.lhf.aicloud.controller;

import com.alibaba.cloud.ai.tongyi.chat.TongYiChatClient;
import com.alibaba.cloud.ai.tongyi.image.TongYiImagesClient;
import com.alibaba.cloud.ai.tongyi.image.TongYiImagesOptions;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.lhf.aicloud.entity.Answer;
import com.lhf.aicloud.entity.enums.AIModelEnum;
import com.lhf.aicloud.entity.enums.AITypeEnum;
import com.lhf.aicloud.service.IAnswerService;
import com.lhf.aicloud.utils.*;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.image.ImageGeneration;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * @author: lhf
 * @date: Created in 2024/08/10
 * @description:
 */
@Slf4j
@RestController
@RequestMapping("/tongyi")
public class TongYiController {

    @Resource
    private TongYiChatClient tongYiChatClient;

    @Resource
    private TongYiImagesClient tongYiImagesClient;

    @Resource
    private IAnswerService answerService;

    @Resource
    private MinIoUtil minIoUtil;

    @Resource
    private RedisTemplate<Object, Object> redisTemplate;

    /**
     * 聊天
     */
    @RequestMapping("/chat")
    public ResponseEntity chat(String question) {
        if (!StringUtils.hasLength(question.trim())) {
            return ResponseEntity.error(400, "问题不能为空");
        }
        log.info("prompt: {}", question);

        String result = tongYiChatClient.call(new Prompt(question))
                .getResult()
                .getOutput()
                .getContent();
        log.info("result: {}", result);

        // 将结果保存到 数据库
        Answer answer = new Answer();
        answer.setQuestion(question);
        answer.setAnswer(result);
        answer.setModel(AIModelEnum.TONGYI.getValue());
        answer.setCreatetime(LocalDateTime.now());
        answer.setUid(SecurityUtils.getCurrentUser().getUid());
        answer.setType(AITypeEnum.CHAT.getValue());
        boolean addResult = answerService.save(answer);
        if (addResult) {
            return ResponseEntity.success(result);
        }
        return ResponseEntity.error("数据保存失败！");
    }

    /**
     * 绘画
     */
    @RequestMapping("/draw")
    public ResponseEntity draw(String question) {

        if (!StringUtils.hasLength(question.trim())) {
            return ResponseEntity.error(400, "问题不能为空");
        }
        log.info("prompt: {}", question);

        String imageRemoteUrl = tongYiImagesClient.call(new ImagePrompt(question,
                TongYiImagesOptions.builder().withHeight(256)
                        .withWidth(256).build())).getResult().getOutput().getUrl();
        log.info("result: {}", imageRemoteUrl);
        // 将图片数据存储的MinIo中
        String imageUrl = "";
        String imageName = "tongyi-draw-" + UUID.randomUUID().toString().replace("-", "") + ".png";
        try (InputStream inputStream = new ByteArrayInputStream(ImageUtils.getImageBytesFromUrl(imageRemoteUrl))) {
            // 上传到 minio
            imageUrl = minIoUtil.upload("tongyi/" + imageName, inputStream, "image/png");

        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        log.info("imageUrl: {}", imageUrl);

        // 将结果保存到 数据库
        Answer answer = new Answer();
        answer.setQuestion(question);
        answer.setAnswer(imageUrl);
        answer.setModel(AIModelEnum.TONGYI.getValue());
        answer.setCreatetime(LocalDateTime.now());
        answer.setUid(SecurityUtils.getCurrentUser().getUid());
        answer.setType(AITypeEnum.DRAW.getValue());
        boolean addResult = answerService.save(answer);
        if (addResult) {
            return ResponseEntity.success(imageUrl);
        }
        return ResponseEntity.error("数据保存失败！");
    }

    /**
     * 获取通义大模型聊天历史信息
     */
    @RequestMapping("/chatHistory")
    public ResponseEntity chatHistory() {
        // 获取当前用户ID
        Long uid = Objects.requireNonNull(SecurityUtils.getCurrentUser()).getUid();
        String cacheKey = AppVariable.getListCacheKey(uid, AIModelEnum.TONGYI.getValue(), AITypeEnum.CHAT.getValue());
        ArrayList cacheList = (ArrayList)redisTemplate.opsForValue().get(cacheKey);
        if (cacheList != null && !cacheList.isEmpty()) {
            return ResponseEntity.success(cacheList);
        }
        // 构建查询条件
        LambdaQueryWrapper<Answer> answerLambdaQueryWrapper = new LambdaQueryWrapper<Answer>()
                .eq(Answer::getUid, uid)
                .eq(Answer::getType, AITypeEnum.CHAT.getValue())
                .eq(Answer::getModel, AIModelEnum.TONGYI.getValue())
                .orderByDesc(Answer::getCreatetime);
        // 查询
        List<Answer> list = answerService.list(answerLambdaQueryWrapper);
        redisTemplate.opsForValue().set(cacheKey, list, 1, TimeUnit.DAYS);
        return ResponseEntity.success(list);
    }

    /**
     * 获取通义大模型绘画历史信息
     */
    @RequestMapping("/drawHistory")
    public ResponseEntity drawHistory() {
        // 获取当前用户ID
        Long uid = Objects.requireNonNull(SecurityUtils.getCurrentUser()).getUid();
        String cacheKey = AppVariable.getListCacheKey(uid, AIModelEnum.TONGYI.getValue(), AITypeEnum.DRAW.getValue());
        ArrayList cacheList = (ArrayList)redisTemplate.opsForValue().get(cacheKey);
        if (cacheList != null && !cacheList.isEmpty()) {
            return ResponseEntity.success(cacheList);
        }
        // 构建查询条件
        LambdaQueryWrapper<Answer> answerLambdaQueryWrapper = new LambdaQueryWrapper<Answer>()
                .eq(Answer::getUid, uid)
                .eq(Answer::getType, AITypeEnum.DRAW.getValue())
                .eq(Answer::getModel, AIModelEnum.TONGYI.getValue())
                .orderByDesc(Answer::getCreatetime);
        // 查询
        List<Answer> list = answerService.list(answerLambdaQueryWrapper);
        redisTemplate.opsForValue().set(cacheKey, list, 1, TimeUnit.DAYS);
        return ResponseEntity.success(list);
    }
}
